//! # 观测性上下文管理模块
//! 
//! 提供观测性上下文的创建、管理和生命周期控制。
//! 支持上下文的继承、属性管理和自动清理。

use std::collections::HashMap;
use std::time::Instant;
use serde::{Serialize, Deserialize};

/// 观测性上下文
/// 
/// 包含分布式追踪、指标收集和日志记录所需的上下文信息。
/// 支持父子关系和属性继承。
#[derive(Debug, Clone)]
pub struct ObservabilityContext {
    /// 追踪ID - 全局唯一，用于关联整个请求链
    pub trace_id: String,
    
    /// Span ID - 当前操作的唯一标识
    pub span_id: String,
    
    /// 父Span ID - 用于构建调用链关系
    pub parent_span_id: Option<String>,
    
    /// 服务名称
    pub service_name: String,
    
    /// 操作名称
    pub operation_name: String,
    
    /// 开始时间
    pub start_time: Instant,
    
    /// 结束时间
    pub end_time: Option<Instant>,
    
    /// 标签 - 用于过滤和聚合
    pub labels: HashMap<String, String>,
    
    /// 自定义属性 - 业务相关的额外信息
    pub attributes: HashMap<String, serde_json::Value>,
    
    /// 状态
    pub state: ContextState,
}

/// 上下文状态
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ContextState {
    /// 活跃状态
    Active,
    /// 已完成
    Completed,
    /// 发生错误
    Error(String),
    /// 已取消
    Cancelled,
}

impl ObservabilityContext {
    /// 创建新的观测性上下文
    /// 
    /// # 参数
    /// * `service_name` - 服务名称
    /// * `operation_name` - 操作名称
    /// 
    /// # 返回值
    /// 新创建的观测性上下文
    pub fn new(service_name: String, operation_name: String) -> Self {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        
        Self {
            trace_id: format!("trace_{}", now),
            span_id: format!("span_{}", now),
            parent_span_id: None,
            service_name,
            operation_name,
            start_time: Instant::now(),
            end_time: None,
            labels: HashMap::new(),
            attributes: HashMap::new(),
            state: ContextState::Active,
        }
    }

    /// 创建子上下文
    /// 
    /// 继承父上下文的trace_id和labels，创建新的span_id。
    /// 
    /// # 参数
    /// * `operation_name` - 子操作名称
    /// 
    /// # 返回值
    /// 新创建的子上下文
    pub fn create_child(&self, operation_name: String) -> Self {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
            
        Self {
            trace_id: self.trace_id.clone(),
            span_id: format!("span_{}", now),
            parent_span_id: Some(self.span_id.clone()),
            service_name: self.service_name.clone(),
            operation_name,
            start_time: Instant::now(),
            end_time: None,
            labels: self.labels.clone(), // 继承父上下文的标签
            attributes: HashMap::new(),   // 子上下文的属性独立
            state: ContextState::Active,
        }
    }

    /// 添加标签
    /// 
    /// # 参数
    /// * `key` - 标签键
    /// * `value` - 标签值
    pub fn add_label(&mut self, key: String, value: String) {
        self.labels.insert(key, value);
    }

    /// 批量添加标签
    /// 
    /// # 参数
    /// * `labels` - 标签映射
    pub fn add_labels(&mut self, labels: HashMap<String, String>) {
        self.labels.extend(labels);
    }

    /// 添加属性
    /// 
    /// # 参数
    /// * `key` - 属性键
    /// * `value` - 属性值
    pub fn add_attribute(&mut self, key: String, value: serde_json::Value) {
        self.attributes.insert(key, value);
    }

    /// 批量添加属性
    /// 
    /// # 参数
    /// * `attributes` - 属性映射
    pub fn add_attributes(&mut self, attributes: HashMap<String, serde_json::Value>) {
        self.attributes.extend(attributes);
    }

    /// 获取标签值
    /// 
    /// # 参数
    /// * `key` - 标签键
    /// 
    /// # 返回值
    /// 标签值的引用，如果不存在则返回None
    pub fn get_label(&self, key: &str) -> Option<&String> {
        self.labels.get(key)
    }

    /// 获取属性值
    /// 
    /// # 参数
    /// * `key` - 属性键
    /// 
    /// # 返回值
    /// 属性值的引用，如果不存在则返回None
    pub fn get_attribute(&self, key: &str) -> Option<&serde_json::Value> {
        self.attributes.get(key)
    }

    /// 标记上下文完成
    pub fn complete(&mut self) {
        self.end_time = Some(Instant::now());
        self.state = ContextState::Completed;
    }

    /// 标记上下文发生错误
    /// 
    /// # 参数
    /// * `error` - 错误信息
    pub fn mark_error(&mut self, error: String) {
        self.end_time = Some(Instant::now());
        self.state = ContextState::Error(error);
    }

    /// 取消上下文
    pub fn cancel(&mut self) {
        self.end_time = Some(Instant::now());
        self.state = ContextState::Cancelled;
    }

    /// 获取持续时间
    /// 
    /// # 返回值
    /// 上下文的持续时间（毫秒），如果未结束则返回None
    pub fn duration_ms(&self) -> Option<u64> {
        self.end_time.map(|end| {
            end.duration_since(self.start_time).as_millis() as u64
        })
    }

    /// 获取当前持续时间
    /// 
    /// # 返回值
    /// 从开始到现在的持续时间（毫秒）
    pub fn elapsed_ms(&self) -> u64 {
        self.start_time.elapsed().as_millis() as u64
    }

    /// 检查上下文是否活跃
    /// 
    /// # 返回值
    /// 如果上下文处于活跃状态返回true，否则返回false
    pub fn is_active(&self) -> bool {
        matches!(self.state, ContextState::Active)
    }

    /// 检查上下文是否已完成
    /// 
    /// # 返回值
    /// 如果上下文已完成返回true，否则返回false
    pub fn is_completed(&self) -> bool {
        matches!(self.state, ContextState::Completed)
    }

    /// 检查上下文是否有错误
    /// 
    /// # 返回值
    /// 如果上下文有错误返回true，否则返回false
    pub fn has_error(&self) -> bool {
        matches!(self.state, ContextState::Error(_))
    }

    /// 获取错误信息
    /// 
    /// # 返回值
    /// 如果上下文有错误返回错误信息，否则返回None
    pub fn get_error(&self) -> Option<&String> {
        match &self.state {
            ContextState::Error(error) => Some(error),
            _ => None,
        }
    }

    /// 克隆基本信息创建新上下文
    /// 
    /// 保留trace_id、service_name和labels，但创建新的span_id和操作名称。
    /// 
    /// # 参数
    /// * `operation_name` - 新操作名称
    /// 
    /// # 返回值
    /// 新的上下文实例
    pub fn fork(&self, operation_name: String) -> Self {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
            
        Self {
            trace_id: self.trace_id.clone(),
            span_id: format!("span_{}", now),
            parent_span_id: None, // fork的上下文不设置父span
            service_name: self.service_name.clone(),
            operation_name,
            start_time: Instant::now(),
            end_time: None,
            labels: self.labels.clone(),
            attributes: HashMap::new(),
            state: ContextState::Active,
        }
    }

    /// 序列化为JSON字符串（用于传输）
    /// 
    /// # 返回值
    /// 序列化后的JSON字符串，出错时返回错误
    pub fn to_json(&self) -> Result<String, serde_json::Error> {
        let serializable = SerializableContext {
            trace_id: &self.trace_id,
            span_id: &self.span_id,
            parent_span_id: self.parent_span_id.as_ref(),
            service_name: &self.service_name,
            operation_name: &self.operation_name,
            labels: &self.labels,
            attributes: &self.attributes,
        };
        serde_json::to_string(&serializable)
    }

    /// 从JSON字符串反序列化（用于接收）
    /// 
    /// # 参数
    /// * `json` - JSON字符串
    /// 
    /// # 返回值
    /// 反序列化后的上下文，出错时返回错误
    pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
        let serializable: SerializableContext = serde_json::from_str(json)?;
        Ok(Self {
            trace_id: serializable.trace_id.to_string(),
            span_id: serializable.span_id.to_string(),
            parent_span_id: serializable.parent_span_id.map(|s| s.to_string()),
            service_name: serializable.service_name.to_string(),
            operation_name: serializable.operation_name.to_string(),
            start_time: Instant::now(), // 重新开始计时
            end_time: None,
            labels: serializable.labels.clone(),
            attributes: serializable.attributes.clone(),
            state: ContextState::Active,
        })
    }
}

/// 可序列化的上下文结构（用于网络传输）
#[derive(Serialize, Deserialize)]
struct SerializableContext<'a> {
    trace_id: &'a str,
    span_id: &'a str,
    parent_span_id: Option<&'a str>,
    service_name: &'a str,
    operation_name: &'a str,
    labels: &'a HashMap<String, String>,
    attributes: &'a HashMap<String, serde_json::Value>,
}

/// 上下文管理器
/// 
/// 负责管理活跃的观测性上下文，提供清理和统计功能。
pub struct ContextManager {
    /// 活跃上下文存储
    contexts: std::sync::Arc<tokio::sync::RwLock<HashMap<String, ObservabilityContext>>>,
}

impl ContextManager {
    /// 创建新的上下文管理器
    pub fn new() -> Self {
        Self {
            contexts: std::sync::Arc::new(tokio::sync::RwLock::new(HashMap::new())),
        }
    }

    /// 存储上下文
    /// 
    /// # 参数
    /// * `context_id` - 上下文ID
    /// * `context` - 观测性上下文
    pub async fn store_context(&self, context_id: String, context: ObservabilityContext) {
        let mut contexts = self.contexts.write().await;
        contexts.insert(context_id, context);
    }

    /// 获取上下文
    /// 
    /// # 参数
    /// * `context_id` - 上下文ID
    /// 
    /// # 返回值
    /// 上下文的副本，如果不存在则返回None
    pub async fn get_context(&self, context_id: &str) -> Option<ObservabilityContext> {
        let contexts = self.contexts.read().await;
        contexts.get(context_id).cloned()
    }

    /// 移除上下文
    /// 
    /// # 参数
    /// * `context_id` - 上下文ID
    /// 
    /// # 返回值
    /// 被移除的上下文，如果不存在则返回None
    pub async fn remove_context(&self, context_id: &str) -> Option<ObservabilityContext> {
        let mut contexts = self.contexts.write().await;
        contexts.remove(context_id)
    }

    /// 更新上下文
    /// 
    /// # 参数
    /// * `context_id` - 上下文ID
    /// * `updater` - 更新函数
    pub async fn update_context<F>(&self, context_id: &str, updater: F) -> bool
    where
        F: FnOnce(&mut ObservabilityContext),
    {
        let mut contexts = self.contexts.write().await;
        if let Some(context) = contexts.get_mut(context_id) {
            updater(context);
            true
        } else {
            false
        }
    }

    /// 获取活跃上下文数量
    /// 
    /// # 返回值
    /// 当前活跃的上下文数量
    pub async fn active_count(&self) -> usize {
        let contexts = self.contexts.read().await;
        contexts.len()
    }

    /// 清理过期上下文
    /// 
    /// # 参数
    /// * `max_age` - 最大存活时间
    /// 
    /// # 返回值
    /// 被清理的上下文数量
    pub async fn cleanup_expired(&self, max_age: std::time::Duration) -> usize {
        let mut contexts = self.contexts.write().await;
        let before_count = contexts.len();
        
        let now = Instant::now();
        contexts.retain(|_, context| {
            now.duration_since(context.start_time) < max_age
        });
        
        let cleaned_count = before_count - contexts.len();
        if cleaned_count > 0 {
            tracing::info!("清理了 {} 个过期的观测性上下文", cleaned_count);
        }
        
        cleaned_count
    }

    /// 获取所有上下文的统计信息
    /// 
    /// # 返回值
    /// 上下文统计信息
    pub async fn get_stats(&self) -> ContextStats {
        let contexts = self.contexts.read().await;
        let mut stats = ContextStats::default();
        
        for context in contexts.values() {
            stats.total_count += 1;
            match &context.state {
                ContextState::Active => stats.active_count += 1,
                ContextState::Completed => stats.completed_count += 1,
                ContextState::Error(_) => stats.error_count += 1,
                ContextState::Cancelled => stats.cancelled_count += 1,
            }
        }
        
        stats
    }
}

impl Default for ContextManager {
    fn default() -> Self {
        Self::new()
    }
}

/// 上下文统计信息
#[derive(Debug, Clone, Default)]
pub struct ContextStats {
    /// 总上下文数量
    pub total_count: usize,
    /// 活跃上下文数量
    pub active_count: usize,
    /// 已完成上下文数量
    pub completed_count: usize,
    /// 错误上下文数量
    pub error_count: usize,
    /// 已取消上下文数量
    pub cancelled_count: usize,
}